import os
import shutil

from tqdm import tqdm


def getFiles(path):
    ans = []
    for root, dirs, files in os.walk(path):
        for i in files:
            ans.append(root + '/' + i)
    return ans


def main():
    RAW_paths = [r'D:\RAWs', r'G:\图\有导出的RAW备份_2', 'E:\\']
    selected_image_paths = [r'G:\图\旧导出', r'D:\图\光影世界']
    valid_extensions = set(['.jpg', '.jpeg', '.png', '.cr2', '.cr3'])
    baked_path = r'G:\图\有导出的RAW备份'

    RAW_files = []
    for i in RAW_paths:
        RAW_files += getFiles(i)

    selected_image_files = []
    for selected_image_path in selected_image_paths:
        selected_image_files.extend(getFiles(selected_image_path))
    selected_image_file_names = [os.path.basename(i) for i in selected_image_files]
    selected_image_file_name_without_exts = [os.path.splitext(i)[0] for i in selected_image_file_names]
    selected_image_file_name_without_exts = set(selected_image_file_name_without_exts)
    # print('selected_image_file_name_without_exts', selected_image_file_name_without_exts)

    baked_file_names = set()

    with tqdm(total=len(selected_image_file_name_without_exts)) as pbar:
        for selected_image_file_name_without_ext in selected_image_file_name_without_exts:
            for i in RAW_files:
                file_name = os.path.basename(i)
                file_name_without_ext, file_ext = os.path.splitext(file_name)
                if file_ext.lower() not in valid_extensions:
                    continue
                if file_name_without_ext in selected_image_file_name_without_ext:
                    pbar.set_description('备份 %s' % (selected_image_file_name_without_ext, ))
                    target_file_path, target_file_path_dir = None, None
                    # 除去RAW_paths的路径
                    for RAW_path in RAW_paths:
                        if i.startswith(RAW_path):
                            dir_path = os.path.dirname(i)
                            file_path_without_prefix = dir_path[len(RAW_path):].strip('\\').strip('/')
                            # print('file_path_without_prefix', file_path_without_prefix)
                            target_file_path_dir = os.path.join(baked_path, file_path_without_prefix)
                            target_file_path = os.path.join(target_file_path_dir, file_name)
                            # print('target_file_path', target_file_path)
                    try:
                        assert target_file_path is not None
                        assert target_file_path_dir is not None
                    except:
                        print('target_file_path is None')
                        print('i', i)
                        print('file_name', file_name)
                        print('file_name_without_ext', file_name_without_ext)
                        print('file_ext', file_ext)
                        raise Exception('target_file_path is None')
                    if not os.path.exists(target_file_path_dir):
                        os.makedirs(target_file_path_dir)
                    if not os.path.exists(target_file_path):
                        shutil.copy(i, target_file_path)
                    baked_file_names.add(selected_image_file_name_without_ext)
            pbar.update(1)
    
    # 没有找到源文件的图片
    if len(selected_image_file_name_without_exts) != len(baked_file_names):
        print('没有找到源文件的图片:')
        for i in selected_image_file_name_without_exts:
            if i not in baked_file_names:
                print(i)


if __name__ == '__main__':
    main()
